/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv;

import com.webank.tikv.operation.ErrorHandler;
import com.webank.tikv.common.RetryPolicy;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.ChannelFactory;
import io.grpc.CallOptions;
import io.grpc.MethodDescriptor;
import io.grpc.stub.AbstractFutureStub;
import io.grpc.stub.AbstractStub;
import io.grpc.stub.ClientCalls;
import io.grpc.stub.StreamObserver;

import java.util.concurrent.Executor;
import java.util.function.Supplier;

public abstract class AbstractGRPCClient<
    BlockingStubT extends AbstractStub<BlockingStubT>,
    FutureStubT extends AbstractFutureStub<FutureStubT>> {

    protected final ChannelFactory channelFactory;
    protected TiConfiguration conf;
    protected long timeout;
    protected BlockingStubT blockingStub;
    protected FutureStubT asyncStub;

    protected AbstractGRPCClient(TiConfiguration conf, ChannelFactory channelFactory) {
        this.conf = conf;
        this.timeout = conf.getTimeout();
        this.channelFactory = channelFactory;
    }

    protected AbstractGRPCClient(TiConfiguration conf, ChannelFactory channelFactory,
                                 BlockingStubT blockingStub, FutureStubT asyncStub) {
        this.conf = conf;
        this.timeout = conf.getTimeout();
        this.channelFactory = channelFactory;
        this.blockingStub = blockingStub;
        this.asyncStub = asyncStub;
    }

    public <ReqT, RespT> RespT callWithRetry(BackOffer backOffer, MethodDescriptor<ReqT, RespT> method,
                                             Supplier<ReqT> requestFactory, ErrorHandler<RespT> handler) {
        RetryPolicy<RespT> policy = new RetryPolicy<>(backOffer, handler);
        return policy.callWithRetry(() -> {
            BlockingStubT stub = getBlockingStub();
            return ClientCalls.blockingUnaryCall(
                stub.getChannel(), method, stub.getCallOptions(), requestFactory.get());
        });
    }

    protected <ReqT, RespT> void callAsyncWithRetry(BackOffer backOffer, Executor executor,
                                                    MethodDescriptor<ReqT, RespT> method,
                                                    Supplier<ReqT> requestFactory,
                                                    StreamObserver<RespT> responseObserver,
                                                    ErrorHandler<RespT> errorhandler) {
        RetryPolicy<RespT> policy = new RetryPolicy<>(backOffer, errorhandler);
        policy.callWithRetry(() -> {
            FutureStubT stub = getAsyncStub();
            CallOptions callOptions = stub.getCallOptions()
                .withExecutor(executor);
            ClientCalls.asyncUnaryCall(
                stub.getChannel().newCall(method, callOptions),
                requestFactory.get(),
                responseObserver);
            return null;
        });
    }

    public void setTimeout(long timeout) {
        this.timeout = timeout;
    }

    public long getTimeout() {
        return timeout;
    }

    protected abstract BlockingStubT getBlockingStub();

    protected abstract FutureStubT getAsyncStub();

}
